"""
简单智能体示例
演示如何使用SDK创建一个基本的智能体
"""
import json
from agent_sdk import AgentSDK


class SimpleAgent:
    def __init__(self):
        self.sdk = AgentSDK(
            platform_url="http://localhost:8080",
            agent_config={
                "name": "SimpleAgent",
                "type": "ANALYSIS",
                "description": "一个简单的数据分析智能体",
                "capabilities": ["data_analysis", "statistics"],
                "endpoint": "http://localhost:8001"
            }
        )
    
    def start(self):
        """启动智能体"""
        # 注册到平台
        result = self.sdk.register()
        print(f"注册成功: {result}")
        
        # 启动心跳
        self.sdk.start_heartbeat(interval=30)
        
        # 设置任务处理回调
        self.sdk.receive_task(self.handle_task)
        
        print("智能体已启动，等待任务...")
        
        # 保持运行
        try:
            import time
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            print("\n正在停止...")
            self.sdk.unregister()
    
    def handle_task(self, task):
        """处理任务"""
        print(f"收到任务: {task}")
        
        try:
            # 解析输入数据
            input_data = json.loads(task.get("inputData", "{}"))
            
            # 执行业务逻辑
            result = self.process(input_data)
            
            # 提交成功结果
            self.sdk.submit_result(
                task_id=task["taskId"],
                status="SUCCESS",
                output_data=result
            )
            print(f"任务完成: {task['taskId']}")
            
        except Exception as e:
            # 提交失败结果
            self.sdk.submit_result(
                task_id=task["taskId"],
                status="FAILED",
                error_message=str(e)
            )
            print(f"任务失败: {e}")
    
    def process(self, input_data):
        """业务逻辑处理"""
        # 这里实现具体的业务逻辑
        print(f"处理数据: {input_data}")
        
        # 示例：简单的数据统计
        if "numbers" in input_data:
            numbers = input_data["numbers"]
            result = {
                "count": len(numbers),
                "sum": sum(numbers),
                "avg": sum(numbers) / len(numbers) if numbers else 0
            }
            return result
        
        return {"status": "processed"}


if __name__ == "__main__":
    agent = SimpleAgent()
    agent.start()
